% Setup simulation parameters.
clc; clear all;

simPar = setupSimParameters();
simPar.v = 20;
simPar.N = 32;
N_sim = 5;
N_scan = simPar.T_pre;

SNR_dB_arr = 5; 
% SNR_dB_arr = (linspace(-10, 15, N_scan).');
T_arr = (linspace(0, simPar.Ts * (simPar.T_pre - 1), simPar.T_pre).') * 1000;

fprintf('Time Range = ');
for idx = 1:N_scan
    fprintf('%.2f ', T_arr(idx));
end
fprintf(' (ms)\n');
fprintf('Start simulation...\n');

% 初始化 NMSE 变量
NMSE_LS = zeros(N_sim, N_scan);
NMSE_MMSE_ISO = zeros(N_sim, N_scan);
NMSE_AR = zeros(N_sim, N_scan);
NMSE_Prony = zeros(N_sim, N_scan);
NMSE_OMP = zeros(N_sim, N_scan); 
NMSE_singleKernelGPR = zeros(N_sim, N_scan);
NMSE_singleKernelGPRmu = zeros(N_sim, N_scan);
NMSE_mixedKernelGPR = zeros(N_sim, N_scan);

Rate_CSI = zeros(N_sim, N_scan);
Rate_LS = zeros(N_sim, N_scan);
Rate_MMSE_ISO = zeros(N_sim, N_scan);
Rate_AR = zeros(N_sim, N_scan);
Rate_Prony = zeros(N_sim, N_scan);
Rate_OMP = zeros(N_sim, N_scan); 
Rate_singleKernelGPR = zeros(N_sim, N_scan);
Rate_singleKernelGPRmu = zeros(N_sim, N_scan);
Rate_mixedKernelGPR = zeros(N_sim, N_scan);

fprintf('+==========+===================================================================================================+\n');
fprintf('|          |                                              NMSE(dB)                                             |\n');
fprintf('| SNR(dB)  +---------------------------------------------------------------------------------------------------+\n');
fprintf('|          |  LS  | MMSE(ISO) | AR(1) | Prony |  OMP  |  GPR  | GPRmu  |  GPR(2) |  Time Elapsed (s) |\n');
fprintf('+----------+------+-----------+-------+-------+-------+-------+--------+---------+-------------------+\n');

% % 启动并行池并附加必要的文件
% if ~exist('fit_GPR.m', 'file')
%     error('The file fit_GPR.m does not exist in the current directory.');
% end
% % if ~exist('get_C_only.m', 'file')
% %     error('The file get_C_only.m does not exist in the current directory.');
% % end
% pool = gcp('nocreate');
% if isempty(pool)
%     pool = parpool();
% end
% addAttachedFiles(pool, 'fit_GPR.m');
tic;
% 使用 parfor 替换外层循环
results = cell(N_sim, 1);
parfor idx_sim = 1:N_sim
%         rng(idx_sim+1);
        stream = RandStream('mt19937ar', 'Seed', idx_sim);
        RandStream.setGlobalStream(stream);
    % 重新初始化 simPar
    % simPar = setupSimParameters();
    % simPar.v = 20;
    % simPar.SNR_dB = SNR_dB_arr;
    % simPar.sigma2_n = 1 / db2pow(simPar.SNR_dB);
    
%     rng(idx_sim + 1);
    user_theta = -pi/3 + pi/3 * rand(); % cdl not use this user
    user = (25 * rand() + 5) * [sin(user_theta); 0; cos(user_theta)]; 

    % [h0, UE0] = get_channel0(simPar, user);
    [h1] = get_channel1(simPar, user);
    h = zeros(simPar.N, simPar.T_pre);
    y = zeros(simPar.N, simPar.T_pre);
    for idx_pre = 1:simPar.T_pre
        h(:, idx_pre) = h1(:, simPar.T + idx_pre);   
        y(:, idx_pre) = h(:, idx_pre) + sqrt(simPar.sigma2_n) * (randn([simPar.N, 1]) + 1i * randn([simPar.N, 1])) / sqrt(2);
    end
    % y0 = h0 + sqrt(simPar.sigma2_n) * (randn([simPar.N, 1]) + 1i * randn([simPar.N, 1])) / sqrt(2);
    yg = zeros(simPar.N * simPar.T, 1);
    yg_prony = zeros(simPar.N, simPar.T);
    y_ = zeros(simPar.N, 1);
    
    for idx = 1:simPar.T
        yg((idx - 1) * simPar.N + 1:idx * simPar.N, 1) = h1(:, idx) + sqrt(simPar.sigma2_n) * (randn([simPar.N, 1]) + 1i * randn([simPar.N, 1])) / sqrt(2);
        yg_prony(:, idx) = yg((idx - 1) * simPar.N + 1:idx * simPar.N, 1);
        y_ = y_ + yg((idx - 1) * simPar.N + 1:idx * simPar.N, 1);
    end
    y_ = y_ / simPar.T;

        Rate_CSI_tmp = zeros(1, N_scan);
    for idx = 1:simPar.T_pre
        Rate_CSI_tmp(idx) = calc_achievable_rate(h(:, idx), h(:, idx),simPar.sigma2_n);
    end
    
    h_hat_LS = predict_channel_more_v_20(yg_prony(:, simPar.T), simPar, "LS", user);
    NMSE_LS_tmp = zeros(1, N_scan);
    Rate_LS_tmp = zeros(1, N_scan);
    for idx = 1:simPar.T_pre
        NMSE_LS_tmp(idx) = calc_NMSE(h_hat_LS, h(:, idx));
        Rate_LS_tmp(idx) = calc_achievable_rate(h_hat_LS, h(:, idx),simPar.sigma2_n);
    end

    h_hat_MMSE_ISO = predict_channel_more_v_20(yg, simPar, "MMSE_ISO", user);
    NMSE_MMSE_ISO_tmp = zeros(1, N_scan);
    Rate_MMSE_ISO_tmp = zeros(1, N_scan);
    for idx = 1:simPar.T_pre
        NMSE_MMSE_ISO_tmp(idx) = calc_NMSE(h_hat_MMSE_ISO((idx - 1) * simPar.N + 1:idx * simPar.N), h(:, idx));
        Rate_MMSE_ISO_tmp(idx) = calc_achievable_rate(h_hat_MMSE_ISO((idx - 1) * simPar.N + 1:idx * simPar.N), h(:, idx),simPar.sigma2_n);
    end
    
    h_AR = yg;
    NMSE_AR_tmp = zeros(1, N_scan);
    Rate_AR_tmp = zeros(1, N_scan);
    for idx = 1:simPar.T_pre
        h_hat_AR = predict_channel_more_v_20(h_AR, simPar, "AR", user);
        NMSE_AR_tmp(idx) = calc_NMSE(h_hat_AR, h(:, idx)); 
        Rate_AR_tmp(idx) = calc_achievable_rate(h_hat_AR, h(:, idx),simPar.sigma2_n);
        h_AR(1:(simPar.T - 1) * simPar.N, 1) = h_AR(simPar.N + 1:simPar.T * simPar.N, 1);
        h_AR((simPar.T - 1) * simPar.N + 1:simPar.T * simPar.N, 1) = h_hat_AR;
    end

    h_Prony = yg_prony;
    h_Prony2 = zeros(simPar.N, simPar.T_pre + simPar.T);
    NMSE_Prony_tmp = zeros(1, N_scan);
    Rate_Prony_tmp = zeros(1, N_scan);
    p = pinv(yg_prony(:,1:simPar.T-1))*yg_prony(:,simPar.T);
    H_RESULT = zeros(simPar.N, simPar.T_pre);
    for idx = 1:simPar.T_pre
        h_hat_Prony = predict_channel_more(h_Prony, simPar, "Prony", user);
        NMSE_Prony_tmp(idx) = calc_NMSE(h_hat_Prony, h(:, idx));
        Rate_Prony_tmp(idx) = calc_achievable_rate(h_hat_Prony, h(:, idx),simPar.sigma2_n);
        h_Prony2(:, 1:idx + 1) = h_Prony;
        h_Prony2(:, idx + 2) = h_hat_Prony;
        h_Prony = h_Prony2(:, 1:idx + simPar.T);
    end
    NMSE_OMP_tmp = zeros(1, N_scan);
    Rate_OMP_tmp = zeros(1, N_scan);
    for idx = 1:simPar.T_pre
        h_hat_OMP = predict_channel_more_v_20(yg_prony(:, simPar.T), simPar, "OMP", user);
        NMSE_OMP_tmp(idx) = calc_NMSE(h_hat_OMP, h(:, idx));
        Rate_OMP_tmp(idx) = calc_achievable_rate(h_hat_OMP, h(:, idx),simPar.sigma2_n);
    end
    
    inner_simPar = simPar; 
    inner_simPar.channelEstimator.GPR.mixedKernel = false;
    h_hat_singleKernelGPR = predict_channel_more_v_20(yg, inner_simPar, "GPR", user);
    NMSE_singleKernelGPR_tmp = zeros(1, N_scan);
    Rate_singleKernelGPR_tmp = zeros(1, N_scan);
    for idx = 1:simPar.T_pre
        NMSE_singleKernelGPR_tmp(idx) = calc_NMSE(h_hat_singleKernelGPR((idx - 1) * simPar.N + 1:idx * simPar.N), h(:, idx));
        Rate_singleKernelGPR_tmp(idx) = calc_achievable_rate(h_hat_singleKernelGPR((idx - 1) * simPar.N + 1:idx * simPar.N), h(:, idx),simPar.sigma2_n);
    end

    h_hat_singleKernelGPRmu = predict_channel_more_v_20(yg, inner_simPar, "GPRmu", user);
    NMSE_singleKernelGPRmu_tmp = zeros(1, N_scan);
    Rate_singleKernelGPRmu_tmp = zeros(1, N_scan);
    for idx = 1:simPar.T_pre
        NMSE_singleKernelGPRmu_tmp(idx) = calc_NMSE(h_hat_singleKernelGPRmu((idx - 1) * simPar.N + 1:idx * simPar.N), h(:, idx));
    end

    inner_simPar.channelEstimator.GPR.mixedKernel = true;
    h_hat_mixedKernelGPR = predict_channel_more_v_20(yg, inner_simPar, "GPR", user);
    NMSE_mixedKernelGPR_tmp = zeros(1, N_scan);
    Rate_mixedKernelGPR_tmp = zeros(1, N_scan);
    for idx = 1:simPar.T_pre
        NMSE_mixedKernelGPR_tmp(idx) = calc_NMSE(h_hat_mixedKernelGPR((idx - 1) * simPar.N + 1:idx * simPar.N), h(:, idx));
        Rate_mixedKernelGPR_tmp(idx) = calc_achievable_rate(h_hat_mixedKernelGPR((idx - 1) * simPar.N + 1:idx * simPar.N), h(:, idx),simPar.sigma2_n);
    end

    % 将结果存储到全局变量中
    results{idx_sim} = struct( ...
        'NMSE_LS', NMSE_LS_tmp, ...
        'NMSE_MMSE_ISO', NMSE_MMSE_ISO_tmp, ...
        'NMSE_AR', NMSE_AR_tmp, ...
        'NMSE_Prony', NMSE_Prony_tmp, ...
        'NMSE_OMP', NMSE_OMP_tmp, ...
        'NMSE_singleKernelGPR', NMSE_singleKernelGPR_tmp, ...
        'NMSE_singleKernelGPRmu', NMSE_singleKernelGPRmu_tmp, ...
        'NMSE_mixedKernelGPR', NMSE_mixedKernelGPR_tmp);
    results_Rate{idx_sim} = struct( ...
        'Rate_LS', Rate_LS_tmp, ...
        'Rate_CSI', Rate_CSI_tmp, ...
        'Rate_MMSE_ISO', Rate_MMSE_ISO_tmp, ...
        'Rate_AR', Rate_AR_tmp, ...
        'Rate_Prony', Rate_Prony_tmp, ...
        'Rate_OMP', Rate_OMP_tmp, ...
        'Rate_singleKernelGPR', Rate_singleKernelGPR_tmp, ...
        'Rate_singleKernelGPRmu', Rate_singleKernelGPRmu_tmp, ...
        'Rate_mixedKernelGPR', Rate_mixedKernelGPR_tmp);
end
timeElapsed = toc();
% 计算平均 NMSE 值
NMSE_LS_mean = zeros(1, N_scan);
NMSE_MMSE_ISO_mean = zeros(1, N_scan);
NMSE_AR_mean = zeros(1, N_scan);
NMSE_Prony_mean = zeros(1, N_scan);
NMSE_OMP_mean = zeros(1, N_scan);
NMSE_singleKernelGPR_mean = zeros(1, N_scan);
NMSE_singleKernelGPRmu_mean = zeros(1, N_scan);
NMSE_mixedKernelGPR_mean = zeros(1, N_scan);

Rate_CSI_mean = zeros(1, N_scan);
Rate_LS_mean = zeros(1, N_scan);
Rate_MMSE_ISO_mean = zeros(1, N_scan);
Rate_AR_mean = zeros(1, N_scan);
Rate_Prony_mean = zeros(1, N_scan);
Rate_OMP_mean = zeros(1, N_scan);
Rate_singleKernelGPR_mean = zeros(1, N_scan);
Rate_singleKernelGPRmu_mean = zeros(1, N_scan);
Rate_mixedKernelGPR_mean = zeros(1, N_scan);

for idx_sim = 1:N_sim
    NMSE_LS_mean = NMSE_LS_mean + results{idx_sim}.NMSE_LS;
    NMSE_MMSE_ISO_mean = NMSE_MMSE_ISO_mean + results{idx_sim}.NMSE_MMSE_ISO;
    NMSE_AR_mean = NMSE_AR_mean + results{idx_sim}.NMSE_AR;
    NMSE_Prony_mean = NMSE_Prony_mean + results{idx_sim}.NMSE_Prony;
    NMSE_OMP_mean = NMSE_OMP_mean + results{idx_sim}.NMSE_OMP;
    NMSE_singleKernelGPR_mean = NMSE_singleKernelGPR_mean + results{idx_sim}.NMSE_singleKernelGPR;
    NMSE_singleKernelGPRmu_mean = NMSE_singleKernelGPRmu_mean + results{idx_sim}.NMSE_singleKernelGPRmu;
    NMSE_mixedKernelGPR_mean = NMSE_mixedKernelGPR_mean + results{idx_sim}.NMSE_mixedKernelGPR;
        Rate_LS_mean = Rate_LS_mean + results_Rate{idx_sim}.Rate_LS;
    Rate_CSI_mean = Rate_CSI_mean + results_Rate{idx_sim}.Rate_CSI;
    Rate_MMSE_ISO_mean = Rate_MMSE_ISO_mean + results_Rate{idx_sim}.Rate_MMSE_ISO;
    Rate_AR_mean = Rate_AR_mean + results_Rate{idx_sim}.Rate_AR;
    Rate_Prony_mean = Rate_Prony_mean + results_Rate{idx_sim}.Rate_Prony;
    Rate_OMP_mean = Rate_OMP_mean + results_Rate{idx_sim}.Rate_OMP;
    Rate_singleKernelGPR_mean = Rate_singleKernelGPR_mean + results_Rate{idx_sim}.Rate_singleKernelGPR;
    Rate_singleKernelGPRmu_mean = Rate_singleKernelGPRmu_mean + results_Rate{idx_sim}.Rate_singleKernelGPRmu;
    Rate_mixedKernelGPR_mean = Rate_mixedKernelGPR_mean + results_Rate{idx_sim}.Rate_mixedKernelGPR;
end

NMSE_LS_mean = NMSE_LS_mean / N_sim;
NMSE_MMSE_ISO_mean = NMSE_MMSE_ISO_mean / N_sim;
NMSE_AR_mean = NMSE_AR_mean / N_sim;
NMSE_Prony_mean = NMSE_Prony_mean / N_sim;
NMSE_OMP_mean = NMSE_OMP_mean / N_sim;
NMSE_singleKernelGPR_mean = NMSE_singleKernelGPR_mean / N_sim;
NMSE_singleKernelGPRmu_mean = NMSE_singleKernelGPRmu_mean / N_sim;
NMSE_mixedKernelGPR_mean = NMSE_mixedKernelGPR_mean / N_sim;

Rate_LS_mean = Rate_LS_mean / N_sim;
Rate_CSI_mean = Rate_CSI_mean / N_sim;
Rate_MMSE_ISO_mean = Rate_MMSE_ISO_mean / N_sim;
Rate_AR_mean = Rate_AR_mean / N_sim;
Rate_Prony_mean = Rate_Prony_mean / N_sim;
Rate_OMP_mean = Rate_OMP_mean / N_sim;
Rate_singleKernelGPR_mean = Rate_singleKernelGPR_mean / N_sim;
Rate_singleKernelGPRmu_mean = Rate_singleKernelGPRmu_mean / N_sim;
Rate_mixedKernelGPR_mean = Rate_mixedKernelGPR_mean / N_sim;

% 打印结果
for idx_scan = 1:N_scan
    fprintf('|  % 6.3f  |% 6.2f|   % 6.2f  | % 6.2f| % 6.2f| % 6.2f| % 6.2f|  % 6.2f|  %6.2f |      %7.2f      |\n', ...
        T_arr(idx_scan), ...
        pow2db(NMSE_LS_mean(idx_scan)), ...
        pow2db(NMSE_MMSE_ISO_mean(idx_scan)), ...
        pow2db(NMSE_AR_mean(idx_scan)), ...
        pow2db(NMSE_Prony_mean(idx_scan)), ...
        pow2db(NMSE_OMP_mean(idx_scan)), ...
        pow2db(NMSE_singleKernelGPR_mean(idx_scan)), ...
        pow2db(NMSE_singleKernelGPRmu_mean(idx_scan)), ...
        pow2db(NMSE_mixedKernelGPR_mean(idx_scan)), ...
        timeElapsed); 
end
fprintf('+==========+===================================================================================================+\n');
fprintf('+==========+===================================================================================================+\n');
fprintf('|          |                                Achievable rate(bit/s/Hz)                                          |\n');
fprintf('| Time(ms)  +---------------------------------------------------------------------------------------------------+\n');
fprintf('|          |  CSI | MMSE(ISO) | AR(1) | Prony |  OMP  |  GPR  | GPRmu  |  GPR(2) |  Time Elapsed (s) |\n');
fprintf('+----------+------+-----------+-------+-------+-------+-------+--------+---------+-------------------+\n');
for idx_scan = 1:N_scan
    fprintf('|  % 6.3f  |% 6.2f|   % 6.2f  | % 6.2f| % 6.2f| % 6.2f| % 6.2f|  % 6.2f|  %6.2f |      %7.2f      |\n', ...
        T_arr(idx_scan), ...
        (Rate_CSI_mean(idx_scan)), ...
        (Rate_MMSE_ISO_mean(idx_scan)), ...
        (Rate_AR_mean(idx_scan)), ...
        (Rate_Prony_mean(idx_scan)), ...
        (Rate_OMP_mean(idx_scan)), ...
        (Rate_singleKernelGPR_mean(idx_scan)), ...
        (Rate_singleKernelGPRmu_mean(idx_scan)), ...
        (Rate_mixedKernelGPR_mean(idx_scan)), ...
        timeElapsed); 
end
fprintf('+==========+===================================================================================================+\n');
fprintf('Sim complete...\n');

%% Visualization.
set(0,'DefaultLineMarkerSize',  4);
set(0,'DefaultTextFontSize',    14);
set(0,'DefaultAxesFontSize',    12);
set(0,'DefaultLineLineWidth',   1.4);
set(0,'defaultfigurecolor',     'w');
figure('color', [1 1 1]);
% 颜色矩阵
C = [   0.0000    0.0000    0.0000
        0.2000    0.4000    0.8494
        0.3718    0.7176    0.4000
        1.0000    0.5482    0.1000
        0.9650    0.1110    0.1330
        0.8650    0.1110    0.8330
        0.0000    0.0000    0.9000
    ]; 
% 绘制每条曲线
p1 = plot(T_arr, pow2db(NMSE_LS_mean),       'Marker', 'd', 'MarkerSize', 6, 'MarkerFaceColor','w', 'Color', C(1, :), 'DisplayName', 'no prediction LS');
hold on;
p2 = plot(T_arr, pow2db(NMSE_MMSE_ISO_mean), 'Marker', 's', 'MarkerSize', 6, 'MarkerFaceColor','w', 'Color', C(3, :), 'DisplayName', 'no learning EM');
p3 = plot(T_arr, pow2db(NMSE_AR_mean),       'Marker', 'o', 'MarkerSize', 6, 'MarkerFaceColor','w', 'Color', C(4, :), 'DisplayName', 'AR(5)');
p4 = plot(T_arr, pow2db(NMSE_Prony_mean),    'Marker', 'o', 'MarkerSize', 6, 'MarkerFaceColor','w', 'Color', C(7, :), 'DisplayName', 'Prony');
p6 = plot(T_arr, pow2db(NMSE_singleKernelGPR_mean), 'Marker', '>', 'MarkerSize', 6, 'MarkerFaceColor','w', 'Color', C(5, :), 'DisplayName', 'EM kernel learning');
p7 = plot(T_arr, pow2db(NMSE_mixedKernelGPR_mean), 'Marker', '>', 'MarkerSize', 6, 'MarkerFaceColor','w', 'Color', C(6, :), 'DisplayName', 'GEM kernel learning');
% 添加网格、标签和图例
grid on;
xlabel('time (s)', 'Interpreter','latex');
ylabel('NMSE (dB)', 'Interpreter','latex');
legend([p1, p2, p3, p4, p6, p7], 'Location', 'best'); % 使用每个 plot 返回的句柄
set(gca,'FontName','Times New Roman');


figure('color', [1 1 1]);
% 颜色矩阵
C = [   0.0000    0.0000    0.0000
        0.2000    0.4000    0.8494
        0.3718    0.7176    0.4000
        1.0000    0.5482    0.1000
        0.9650    0.1110    0.1330
        0.8650    0.1110    0.8330
        0.0000    0.0000    0.9000
    ]; 
% 绘制每条曲线
% p1 = plot(T_arr, (Rate_LS_mean),       'Marker', 'd', 'MarkerSize', 6, 'MarkerFaceColor','w', 'Color', C(1, :), 'DisplayName', 'no prediction LS');
p1 = plot(T_arr, (Rate_CSI_mean),       'LineStyle', '-.', 'MarkerSize', 6, 'MarkerFaceColor','w', 'Color', C(1, :), 'DisplayName', 'Perfect CSI');
hold on;
p2 = plot(T_arr, (Rate_MMSE_ISO_mean), 'Marker', 's', 'MarkerSize', 6, 'MarkerFaceColor','w', 'Color', C(3, :), 'DisplayName', 'no learning EM');
p3 = plot(T_arr, (Rate_AR_mean),       'Marker', 'o', 'MarkerSize', 6, 'MarkerFaceColor','w', 'Color', C(4, :), 'DisplayName', 'AR(5)');
p4 = plot(T_arr, (Rate_Prony_mean),    'Marker', 'o', 'MarkerSize', 6, 'MarkerFaceColor','w', 'Color', C(7, :), 'DisplayName', 'Prony');
p6 = plot(T_arr, (Rate_singleKernelGPR_mean), 'Marker', '>', 'MarkerSize', 6, 'MarkerFaceColor','w', 'Color', C(5, :), 'DisplayName', 'EM kernel learning');
p7 = plot(T_arr, (Rate_mixedKernelGPR_mean), 'Marker', '>', 'MarkerSize', 6, 'MarkerFaceColor','w', 'Color', C(6, :), 'DisplayName', 'GEM kernel learning');
% 添加网格、标签和图例
grid on;
xlabel('time (ms)', 'Interpreter','latex');
ylabel('Achievable rate(bit/s/Hz)', 'Interpreter','latex');
legend([p1, p2, p3, p4, p6, p7], 'Location', 'best'); % 使用每个 plot 返回的句柄
set(gca,'FontName','Times New Roman');
% Problem: when mu is large, why is the correlation decaying?
%% Save files
save('results/Ray_moretime_v20_NMSE.mat');

